function out = mu_ic_schedule(H, info, cb, opts)
%MU_IC_SCHEDULE_MT Receiver/GDFE schedule builder for minPICU_mt outputs.
%
%   out = mu_ic_schedule_mt(H, info, cb)
%   out = mu_ic_schedule_mt(H, info, cb, opts)
%
% This is a multitone-compatible replacement for mu_ic_schedule.m.  It keeps
% the old N=1 organization, but for N>1 it builds a tone-indexed schedule:
%
%       out.schedule(s).tone(n).rx(r)
%
% where s is the time-share index, n is the tone, and r is the receiver.
%
% The input minPICU_mt stores per-share atom_energies and atom_rates as
% K x N arrays.  The old one-tone scheduler flattened these arrays and then
% used only atom indices 1:K, which fails as soon as N > 1.

    if nargin < 4
        opts = struct();
    end
    if ~isfield(opts,'activeTol'), opts.activeTol = 1e-7; end
    if ~isfield(opts,'lambdaTol'), opts.lambdaTol = 1e-5; end
    if ~isfield(opts,'verbose'),   opts.verbose   = 1; end

    activeTol = opts.activeTol;
    lambdaTol = opts.lambdaTol;

    if ndims(H) == 2
        H = reshape(H, size(H,1), size(H,2), 1);
    end

    atoms = info.atom_table;
    U = info.U;
    K = info.K;

    if isfield(info,'N') && ~isempty(info.N)
        N = info.N;
    else
        N = size(H,3);
    end

    assert(size(H,1) == U && size(H,2) == U, 'H must be U x U x N.');
    assert(size(H,3) == N, 'size(H,3) must match info.N.');

    out = struct();
    out.U = U;
    out.K = K;
    out.N = N;
    out.cb = cb;
    out.atom_table = atoms;
    out.schedule = info.schedule;

    weighted_user_rates = zeros(U,1);
    weighted_user_energies = zeros(U,1);

    for s = 1:numel(info.schedule)

        Pkn = reshape_kn(info.schedule(s).atom_energies, K, N, 'atom_energies');
        Rkn = reshape_kn(info.schedule(s).atom_rates,    K, N, 'atom_rates');

        out.schedule(s).atom_energies_kn = Pkn;
        out.schedule(s).atom_rates_kn = Rkn;
        out.schedule(s).user_rates_by_tone = user_sums_by_tone(Rkn, atoms, U);
        out.schedule(s).user_energies_by_tone = user_sums_by_tone(Pkn, atoms, U);
        out.schedule(s).user_rates = sum(out.schedule(s).user_rates_by_tone, 2);
        out.schedule(s).user_energies = sum(out.schedule(s).user_energies_by_tone, 2);

        lam = getfield_default(info.schedule(s), 'lambda', 1);
        weighted_user_rates = weighted_user_rates + lam * out.schedule(s).user_rates(:);
        weighted_user_energies = weighted_user_energies + lam * out.schedule(s).user_energies(:);

        % Multitone tensor/cell outputs.  These are indexed by {receiver,tone}
        % to avoid trying to force variable-size GDFE matrices into numeric tensors.
        out.schedule(s).GU       = cell(U,N);
        out.schedule(s).WU       = cell(U,N);
        out.schedule(s).S0       = cell(U,N);
        out.schedule(s).S0diag   = cell(U,N);
        out.schedule(s).MSWMFU   = cell(U,N);
        out.schedule(s).bits     = cell(U,N);
        out.schedule(s).dec      = cell(U,N);
        out.schedule(s).noise    = cell(U,N);
        out.schedule(s).rx       = repmat(empty_rx_struct(), U, N);
        out.schedule(s).tone     = struct([]);

        for n = 1:N
            P = Pkn(:,n);
            R = Rkn(:,n);
            Aatom = diag(sqrt(max(P,0)));

            tone_struct = struct();
            tone_struct.n = n;
            tone_struct.atom_energies = P;
            tone_struct.atom_rates = R;
            tone_struct.user_rates = out.schedule(s).user_rates_by_tone(:,n);
            tone_struct.user_energies = out.schedule(s).user_energies_by_tone(:,n);
            tone_struct.rx = repmat(empty_rx_struct(), U, 1);

            for r = 1:U

                ord0 = info.schedule(s).rx_orders{r};
                ord0 = ord0(:).';

                keep = false(size(ord0));
                for m = 1:numel(ord0)
                    k = ord0(m);
                    keep(m) = (P(k) > activeTol) || (R(k) > activeTol);
                end
                dec = ord0(keep);

                hrow = zeros(1,K);
                for k = 1:K
                    u = atoms(k).user;
                    hrow(k) = H(r,u,n);
                end

                Hr_atom = hrow * Aatom;

                active_atoms = find(P(:).' > activeTol);
                noise = setdiff(active_atoms, dec, 'stable');

                rx = empty_rx_struct();
                rx.receiver = r;
                rx.tone = n;
                rx.decode_order_original = ord0;
                rx.decode_order = dec;
                rx.decode_labels = labels_for_atoms(atoms, dec);

                rx.noise_atoms = noise;
                rx.noise_labels = labels_for_atoms(atoms, noise);

                rx.atom_rates = R(dec);
                rx.atom_energies = P(dec);

                rx.H_atom = Hr_atom;
                rx.H_decode_ordered = Hr_atom(dec);
                rx.H_noise = Hr_atom(noise);

                rx.sinr = zeros(numel(dec),1);
                rx.bits = zeros(numel(dec),1);

                for m = 1:numel(dec)
                    k = dec(m);
                    future = dec(m+1:end);
                    interf = [future noise];

                    denom = 1 + sum(abs(Hr_atom(interf)).^2);
                    numer = abs(Hr_atom(k)).^2;

                    rx.sinr(m) = numer / denom;
                    rx.bits(m) = (1/cb)*log2(1 + rx.sinr(m));
                end

                [rx.GU, rx.WU, rx.S0, rx.MSWMFU, rx.gdfe_bits] = ...
                    scalar_ordered_gdfe(Hr_atom, dec, noise, cb, rx.bits);

                if isempty(rx.S0)
                    rx.S0diag = [];
                else
                    rx.S0diag = diag(rx.S0).';
                end

                out.schedule(s).rx(r,n) = rx;
                tone_struct.rx(r) = rx;

                out.schedule(s).GU{r,n}       = rx.GU;
                out.schedule(s).WU{r,n}       = rx.WU;
                out.schedule(s).S0{r,n}       = rx.S0;
                out.schedule(s).S0diag{r,n}   = rx.S0diag;
                out.schedule(s).MSWMFU{r,n}   = rx.MSWMFU;
                out.schedule(s).bits{r,n}     = rx.gdfe_bits;
                out.schedule(s).dec{r,n}      = dec;
                out.schedule(s).noise{r,n}    = noise;
            end

            out.schedule(s).tone(n) = tone_struct;
        end
    end

    out.weighted_user_rates = weighted_user_rates;
    out.weighted_user_energies = weighted_user_energies;
    if isfield(info,'bu_a')
        out.rate_error_vs_info_bu_a = weighted_user_rates(:) - info.bu_a(:);
    else
        out.rate_error_vs_info_bu_a = [];
    end

    if opts.verbose
        mu_ic_print_schedule_mt(out, info, lambdaTol, activeTol);
    end
end

function X = reshape_kn(Xin, K, N, name)
    if isempty(Xin)
        X = zeros(K,N);
        return;
    end
    if isequal(size(Xin), [K N])
        X = Xin;
    elseif isvector(Xin) && numel(Xin) == K*N
        X = reshape(Xin(:), K, N);
    elseif N == 1 && numel(Xin) == K
        X = Xin(:);
    else
        error('%s must be K x N or a vector of length K*N.  Got %s, expected %d x %d.', ...
              name, mat2str(size(Xin)), K, N);
    end
end

function S = user_sums_by_tone(Xkn, atoms, U)
    N = size(Xkn,2);
    S = zeros(U,N);
    for u = 1:U
        idx = find([atoms.user] == u);
        S(u,:) = sum(Xkn(idx,:), 1);
    end
end

function rx = empty_rx_struct()
    rx = struct( ...
        'receiver', [], ...
        'tone', [], ...
        'decode_order_original', [], ...
        'decode_order', [], ...
        'decode_labels', {{}}, ...
        'noise_atoms', [], ...
        'noise_labels', {{}}, ...
        'atom_rates', [], ...
        'atom_energies', [], ...
        'H_atom', [], ...
        'H_decode_ordered', [], ...
        'H_noise', [], ...
        'sinr', [], ...
        'bits', [], ...
        'GU', [], ...
        'WU', [], ...
        'S0', [], ...
        'MSWMFU', [], ...
        'gdfe_bits', [], ...
        'S0diag', []);
end

function c = labels_for_atoms(atoms, idx)
    if isempty(idx)
        c = {};
    else
        c = {atoms(idx).label};
    end
end

function v = getfield_default(s, fld, defaultVal)
    if isstruct(s) && isfield(s, fld) && ~isempty(s.(fld))
        v = s.(fld);
    else
        v = defaultVal;
    end
end

function [GU, WU, S0, MSWMFU, bits] = scalar_ordered_gdfe(H_atom, dec, noise, cb, sic_bits)

    dec = dec(:).';
    noise = noise(:).';

    M = numel(dec);

    if M == 0
        GU = [];
        WU = [];
        S0 = [];
        MSWMFU = [];
        bits = [];
        return;
    end

    gd = H_atom(dec);
    gn = H_atom(noise);

    N0 = 1 + sum(abs(gn).^2);
    hw = gd / sqrt(N0);

    C = eye(M) + hw' * hw;

    if exist('lohc','file') == 2
        R = lohc(C);
    else
        % MATLAB fallback if the course-reader lohc.m is not on path.
        % chol returns upper R with C = R' * R, so use R' for the monic
        % factorization convention below only as a numerically safe fallback.
        R = chol(C, 'upper');
    end

    s0diag = abs(diag(R)).^2;
    Dhalf = diag(sqrt(s0diag));

    GU = R / Dhalf;
    S0 = diag(s0diag);

    GU = triu(GU);
    GU(1:M+1:end) = 1;

    s0diag = real(diag(S0));
    bits = (1/cb) * log2(s0diag);

    if nargin >= 5 && ~isempty(sic_bits)
        err = max(abs(bits(:) - sic_bits(:)));
        if err > 1e-6
            warning('scalar_ordered_gdfe:bitMismatch', ...
                'GDFE bits and explicit SIC bits differ by %.3e.', err);
        end
    end

    MSWMFU = hw;
    WU = MSWMFU;
end

function mu_ic_print_schedule_mt(out, info, lambdaTol, activeTol)

    fprintf('\n===== FINAL TIME-SHARING RECEIVER SCHEDULE, MULTITONE =====\n');
    fprintf('U = %d, K = %d, N = %d tones\n', out.U, out.K, out.N);

    for s = 1:numel(out.schedule)
        lam = getfield_default(out.schedule(s), 'lambda', 1);
        if abs(lam) < lambdaTol
            continue;
        end

        fprintf('\n============================================================\n');
        fprintf('Share %d: lambda = %.6g', s, lam);
        if isfield(out.schedule(s),'pool_index') && ~isempty(out.schedule(s).pool_index)
            fprintf('  (pool idx %d)', out.schedule(s).pool_index);
        end
        fprintf('\n============================================================\n');

        fprintf('\nShare %d user-rate sums over all tones:\n', s);
        disp(out.schedule(s).user_rates.');
        fprintf('Share %d user-energy sums over all tones:\n', s);
        disp(out.schedule(s).user_energies.');

        fprintf('Per-tone user rates, rows=user and columns=tone:\n');
        disp(out.schedule(s).user_rates_by_tone);
        fprintf('Per-tone user energies, rows=user and columns=tone:\n');
        disp(out.schedule(s).user_energies_by_tone);

        for n = 1:out.N
            fprintf('\n--- Share %d, Tone %d ---\n', s, n);
            for r = 1:out.U
                rx = out.schedule(s).tone(n).rx(r);
                if isempty(rx.decode_order)
                    fprintf('Receiver %d: no active decoded atoms.\n', r);
                    continue;
                end
                fprintf('Receiver %d decoded order: ', r);
                print_atom_list(rx.decode_order, out.atom_table);
                if ~isempty(rx.noise_atoms)
                    fprintf('Receiver %d noise atoms:   ', r);
                    print_atom_list(rx.noise_atoms, out.atom_table);
                end
            end
        end
    end

    fprintf('\nWeighted user-rate vector sum_s lambda_s * atom-rate-sum_s:\n');
    disp(out.weighted_user_rates.');
    fprintf('Weighted user-energy vector sum_s lambda_s * atom-energy-sum_s:\n');
    disp(out.weighted_user_energies.');

    if isfield(info,'bu_a') && ~isempty(info.bu_a)
        fprintf('info.bu_a:\n');
        disp(info.bu_a(:).');
        fprintf('weighted_user_rates - info.bu_a:\n');
        disp(out.weighted_user_rates(:).' - info.bu_a(:).');
    end

    fprintf('Active tolerance = %.3g\n', activeTol);
end

function print_atom_list(idx, atoms)
    for ii = 1:numel(idx)
        k = idx(ii);
        if ii > 1
            fprintf(', ');
        end
        fprintf('%d:%s', k, atoms(k).label);
    end
    fprintf('\n');
end
